import base64
import logging
from typing import Any, Optional, cast

import numpy as np
from sqlalchemy.exc import IntegrityError
from flask import current_app # Added for accessing app config

from core.entities.embedding_type import EmbeddingInputType
from core.model_manager import ModelInstance
from core.model_runtime.entities.model_entities import ModelPropertyKey
from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
from core.rag.embedding.embedding_base import Embeddings
from extensions import db # Changed from extensions.ext_database import db
from extensions import redis_client # Changed from extensions.ext_redis import redis_client. Assuming redis_client is configured and initialized elsewhere if used
from libs import helper
from models.dataset import Embedding # Assuming this is the refactored model

logger = logging.getLogger(__name__)


class CacheEmbedding(Embeddings):
    def __init__(self, model_instance: ModelInstance, user: Optional[str] = None) -> None:
        self._model_instance = model_instance
        self._user = user

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        """Embed search docs in batches of 10."""
        text_embeddings: list[Any] = [None for _ in range(len(texts))]
        embedding_queue_indices = []
        for i, text in enumerate(texts):
            hash_val = helper.generate_text_hash(text) # Renamed variable to avoid conflict with built-in hash
            embedding_record = (
                db.session.query(Embedding)
                .filter_by(
                    model_name=self._model_instance.model, hash=hash_val, provider_name=self._model_instance.provider
                )
                .first()
            )
            if embedding_record:
                text_embeddings[i] = embedding_record.get_embedding()
            else:
                embedding_queue_indices.append(i)
        
        if embedding_queue_indices:
            embedding_queue_texts = [texts[i] for i in embedding_queue_indices]
            embedding_queue_embeddings = []
            try:
                model_type_instance = cast(TextEmbeddingModel, self._model_instance.model_type_instance)
                model_schema = model_type_instance.get_model_schema(
                    self._model_instance.model, self._model_instance.credentials
                )
                max_chunks = (
                    model_schema.model_properties[ModelPropertyKey.MAX_CHUNKS]
                    if model_schema and model_schema.model_properties and ModelPropertyKey.MAX_CHUNKS in model_schema.model_properties
                    else 1
                )
                for i in range(0, len(embedding_queue_texts), max_chunks):
                    batch_texts = embedding_queue_texts[i : i + max_chunks]

                    embedding_result = self._model_instance.invoke_text_embedding(
                        texts=batch_texts, user=self._user, input_type=EmbeddingInputType.DOCUMENT
                    )

                    for vector in embedding_result.embeddings:
                        try:
                            normalized_embedding = (vector / np.linalg.norm(vector)).tolist()
                            if np.isnan(normalized_embedding).any():
                                logger.warning(f"Normalized embedding is nan: {normalized_embedding}")
                                continue
                            embedding_queue_embeddings.append(normalized_embedding)
                        except IntegrityError: # This catch might be too broad here, usually for DB operations
                            db.session.rollback()
                        except Exception:
                            logging.exception("Failed to transform embedding")
                
                cache_hashes = [] # Renamed to avoid conflict
                try:
                    for i_idx, n_embedding in zip(embedding_queue_indices, embedding_queue_embeddings):
                        text_embeddings[i_idx] = n_embedding
                        current_text_hash = helper.generate_text_hash(texts[i_idx])
                        if current_text_hash not in cache_hashes:
                            embedding_cache_obj = Embedding(
                                model_name=self._model_instance.model,
                                hash=current_text_hash,
                                provider_name=self._model_instance.provider,
                            )
                            embedding_cache_obj.set_embedding(n_embedding)
                            db.session.add(embedding_cache_obj)
                            cache_hashes.append(current_text_hash)
                    db.session.commit()
                except IntegrityError:
                    db.session.rollback()
                    logger.warning("Integrity error while saving embeddings to DB cache.")
            except Exception as ex:
                db.session.rollback()
                logger.exception("Failed to embed documents: %s", ex)
                raise ex

        return text_embeddings

    def embed_query(self, text: str) -> list[float]:
        """Embed query text."""
        hash_val = helper.generate_text_hash(text) # Renamed variable
        embedding_cache_key = f"{self._model_instance.provider}_{self._model_instance.model}_{hash_val}"
        
        if redis_client:
            try:
                embedding_data = redis_client.get(embedding_cache_key)
                if embedding_data:
                    redis_client.expire(embedding_cache_key, 600)
                    decoded_embedding = np.frombuffer(base64.b64decode(embedding_data), dtype="float32") # Using float32 as common for embeddings
                    return [float(x) for x in decoded_embedding]
            except Exception as e:
                logger.warning(f"Redis cache read error for query embedding: {e}")

        try:
            embedding_result = self._model_instance.invoke_text_embedding(
                texts=[text], user=self._user, input_type=EmbeddingInputType.QUERY
            )

            embedding_values = embedding_result.embeddings[0]
            embedding_values = (embedding_values / np.linalg.norm(embedding_values)).tolist()
            if np.isnan(embedding_values).any():
                raise ValueError("Normalized embedding is nan please try again")
        except Exception as ex:
            if current_app.config.get("DEBUG", False):
                # Corrected the f-string syntax below
                logging.exception(f"Failed to embed query text '{text[:10]}...({len(text)} chars)'")
            raise ex

        if redis_client:
            try:
                embedding_vector = np.array(embedding_values, dtype="float32")
                vector_bytes = embedding_vector.tobytes()
                encoded_vector = base64.b64encode(vector_bytes)
                encoded_str = encoded_vector.decode("utf-8")
                redis_client.setex(embedding_cache_key, 600, encoded_str)
            except Exception as ex:
                if current_app.config.get("DEBUG", False):
                    # Corrected the f-string syntax below
                    logging.exception(f"Failed to add embedding to redis for the text '{text[:10]}...({len(text)} chars)'")
                # Not re-raising here, as failure to cache shouldn't fail the embedding itself

        return embedding_values

